package com.github.sbb.sherry.primitives.fetcher.mongodb;

import com.github.sbb.sherry.common.base.Preconditions;
import com.github.sbb.sherry.common.base.ValidateResult;
import com.github.sbb.sherry.primitives.config.datasource.DataSource;
import com.github.sbb.sherry.primitives.config.datasource.DataSourceQueryException;
import com.github.sbb.sherry.primitives.config.datasource.MongoDBDataSourceConfig;
import com.mongodb.MongoClientSettings;
import com.mongodb.MongoCredential;
import com.mongodb.ServerAddress;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoClients;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.bson.Document;

public class MongoDBDataSource implements DataSource {

  private MongoClient mongoClient;
  private MongoDBDataSourceConfig config;

  public MongoDBDataSource(MongoDBDataSourceConfig config) {
    Preconditions.checkNotNull(config, "MongoDBDataSourceConfig can`t be null");
    ValidateResult validate = config.validate();
    Preconditions.checkArgument(validate.isPass(), validate.getMessage());

    this.config = config;

    MongoCredential credential = MongoCredential
        .createCredential(config.getUsername(), config.getDatabase(), config.getPassword().toCharArray());

    mongoClient = MongoClients.create(
        MongoClientSettings.builder()
            .applyToClusterSettings(builder ->
                builder.hosts(Collections.singletonList(new ServerAddress(config.getHost(), config.getPort()))))
            .credential(credential)
            .applyToConnectionPoolSettings(builder -> {
              builder.maxSize(config.getMaxSize())
                  .minSize(config.getMinSize())
                  .maxWaitTime(config.getMaxWaitTimeMS(), TimeUnit.MILLISECONDS);
            })
            .build());
  }

  public Map<String, Object> selectOne(String collection, String condition) throws DataSourceQueryException {
    try {
      MongoCollection<Document> mongoCollection = mongoClient.getDatabase(config.getDatabase())
          .getCollection(collection);
      Document conditionDocument = Document.parse(condition);
      FindIterable<Document> documents = mongoCollection.find(conditionDocument).limit(1);
      MongoCursor<Document> iterator = documents.iterator();
      if (iterator.hasNext()) {
        return iterator.next();
      }
      return new HashMap<>();
    } catch (Exception e) {
      throw new DataSourceQueryException(e);
    }
  }

  public List<Map<String, Object>> selectMany(String collection, String condition, int limit)
      throws DataSourceQueryException {
    try {
      MongoCollection<Document> mongoCollection = mongoClient.getDatabase(config.getDatabase())
          .getCollection(collection);
      Document conditionDocument = Document.parse(condition);
      FindIterable<Document> documents = mongoCollection.find(conditionDocument);
      if (limit > 0) {
        documents.limit(limit);
      }
      MongoCursor<Document> iterator = documents.iterator();

      List<Map<String, Object>> result = new ArrayList<>();

      while (iterator.hasNext()) {
        result.add(iterator.next());
      }
      return result;
    } catch (Exception e) {
      throw new DataSourceQueryException(e);
    }
  }

  @Override
  public void check() {
    Preconditions.checkNotNull(mongoClient);
    mongoClient.listDatabases();
  }

  @Override
  public void close() {
    if (mongoClient != null) {
      mongoClient.close();
      mongoClient = null;
    }
  }
}
